#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Mar 21 20:43:36 2019
generate wanted alps lattice
this time Boson ssh hofstadter model ladder lattice
@author: bb
"""

class Vertex(object):
    # a vertex object that contains vertex ID, coordinates
    # type is always 0, coodinates are for ladder
    def __init__(self,ID):
        assert type(ID)==int,'ID has to be an integer.'
        self.ID = ID
        self.type = 0
        if ID%2==0: # if even
            self.coordinate=((ID-2)/2,1)
        else: # if odd
            self.coordinate=((ID-1)/2,0)
    def get_string(self):
        return '<VERTEX id="'+str(self.ID)+'" type="'+str(self.type)+'"><COORDINATE>'\
    +str(self.coordinate[0])+' '+str(self.coordinate[1])+'</COORDINATE></VERTEX>'

class Edge(object):
    Edge_ID = 1 # 代表出生顺序
    # an edge object that contains edge ID, source-target vertex ID, coordinate and type
    # Edge已经改成了general的
    # Edge初始化时根据始末vertex ID数确定内容
    def __init__(self,vertexID1,vertexID2,edgetype,edgevector):
        assert type(vertexID1)==int and type(vertexID2)==int,'ID has to be an integer.'
        assert vertexID1<vertexID2,'vertexID1 should be smaller than vertexID2'
        
        self.ID = Edge.Edge_ID
        self.source = vertexID1
        self.target = vertexID2
        self.type = edgetype
        self.vector = edgevector
        Edge.Edge_ID += 1
        
    def get_string(self):
        return '<EDGE source="'+str(self.source)+'" target="'+str(self.target)+'" id="'\
    +str(self.ID)+'" type="'+str(self.type)+'" vector="'+self.vector+'"/>'

def generate_ssh_ladder_lattice(L):
    assert type(L)==int,'L must be integer.'
    ## ladder的leg数，所以格点数应该翻倍
    IDs = range(1,2*L+1) # ID is from 1,the last one of range is never reached
    vertexes = [Vertex(ID) for ID in IDs] # generate vertexes
    
    edges0 = []  # for K, inter leg edges
    list0 = [i for i in range(1,2*L,2)] 
    edges0 = [Edge(i,i+1,edgetype=0,edgevector='0 1') for i in list0]
    
    list1 = [] # for t1*exp(-1i*phi/2)
    list2 = [] # for t1*exp(1i*phi/2)
    list3 = [] # for t2*exp(-1i*phi/2)
    list4 = [] # for t2*exp(1i*phi/2)
    ID = 1
    for i in range(1,2*L-2,2):
        if ID==1:
            list1.append(i)
            list2.append(i+1)
            ID = 2
        elif ID==2:
            list3.append(i)
            list4.append(i+1)
            ID = 1
    edges1 = [Edge(i,i+2,edgetype=1,edgevector='1 0') for i in list1] # for t1*exp(-1i*phi/2)
    edges2 = [Edge(i,i+2,edgetype=2,edgevector='1 0') for i in list2] # for t1*exp(1i*phi/2)
    edges3 = [Edge(i,i+2,edgetype=3,edgevector='1 0') for i in list3] # for t2*exp(-1i*phi/2)
    edges4 = [Edge(i,i+2,edgetype=4,edgevector='1 0') for i in list4] # for t2*exp(1i*phi/2)
    
    edges = edges0+edges1+edges2+edges3+edges4 #generate edges
    
    ## generate lattice xml file
    f=open('ssh_ladder_lattice_L%s.xml'%(L),'w')
    f.write('<LATTICES>\n')
    f.write('<GRAPH name="ssh_ladder_lattice" dimension="2" vertices="%s" edges="%s">\n'%(len(vertexes),len(edges)))
    for vertex in vertexes:
        f.write(vertex.get_string()+'\n')
    for edge in edges:
        f.write(edge.get_string()+'\n')
    f.write('</GRAPH>\n')
    f.write('</LATTICES>\n')
    f.close()
    
    lattice_library = 'ssh_ladder_lattice_L%s.xml'%(L)
    lattice_name = 'ssh_ladder_lattice'
    return (lattice_library,lattice_name)

#generate_ssh_ladder_lattice(5)